
* Helper program to print results
capture program drop print_results
program define print_results
	syntax varlist ,  model_names( str  ) stat_names(str) stat_labs(str) [outfile(str )  hidevars_print(real 0 ) ]

	if "`outfile'" == "" & `hidevars_print' == 0  {

		#d ;
			esttab `model_names' ,
			style(tab)  nonumbers star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3) labels("`stat_labs'" ))
			drop(o.*, relax) keep(`varlist' , relax) append;
		#d cr
	}
	else if "`outfile'" == "" & `hidevars_print' == 1 {
		#d ;
			esttab `model_names' ,
			style(tab) nonumbers star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3) labels("`stat_labs'" ))
			drop(o.*, relax) keep( , relax) append;
		#d cr
	}
	else if "`outfile'" != "" & `hidevars_print' == 0 {
		#d ;
			esttab `model_names' using "`outfile'",
			style(tab) nonumbers  star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3)  labels("`stat_labs'" ))
			drop(o.*, relax) keep(`varlist' , relax) replace;
		#d cr
	}
	else if "`outfile'" != "" & `hidevars_print' == 1 {
		#d ;
			esttab `model_names' using "`outfile'",
			style(tab) nonumbers  star(* 0.10 ** 0.05 *** 0.01 )
			cells(b(star fmt(%9.3f)) se(par))  collabels(none) depvar
			stats(`stat_names', fmt(a3)  labels("`stat_labs'" ))
			drop(o.*, relax) keep( , relax) replace;
		#d cr
	}

end


* Helper program to combine csvs
capture program drop combine_files
program define combine_files
	syntax varlist ,  outfile(str) file_list(str) suffix(str)

	foreach file in local file_list {
		preserve
		import delimited using "`file'_`suffix'"
		tempfile "`file'_`suffix'"
		save ``file'_`suffix''
		restore
		append using ``file'_`suffix''
	}
	export delimited using "`outfile'", replace
end




* Program to run Balance, First Stage, RF and 2SLS regressions and save results.
capture program drop get_estimates
* need the following globals: balance scores enrollment
program define get_estimates


	syntax , endogvar(varlist) instruments(varlist) idvar(varname)   [ pscore(varlist) outcomes(varlist) balance_vars(varlist) other_controls(varlist fv ) ///
			regression_controls(varlist fv ) outfile(str) balance(real 0 ) cluster(varname) ///
			fs(real 0) 2sls(real 0 ) rf(real 0 ) weight(varlist) hidevars(real 0 ) rf_non_offered_mean(real 0) ///
			non_offered_mean(real 0 )  pointsofsupport(varlist) 2sls_residuals(real 0) mean_offer(real 0) model_file(str)  joint(real 0) restriction(str) ///
			lin_cntrl(real 0) rf_non_offered_mean( real 1)]
	// balance
	local j = 1
	if "`restriction'" != "" {
		keep if `restriction'
	}
	if "`pscore'" != "" {
		keep if !inlist(`pscore',0,1)
	}
	if `balance' == 1 {
		* just for esttab
		local endogbal = ""
		foreach var of local instruments{

			gen `var'_bl = `var'
			local endogbal " `endogbal' `var'_bl "

		}

		local estimatesstring ""

		foreach y of local balance_vars {

			if "`weight'" != ""{
				qui:  reg `y'  `endogbal' `other_controls' `fvars' [aw=`weight'], robust
			}
			else if "`cluster'" != "" {
				qui: reg `y'  `endogbal' `other_controls' `fvars' , vce(cluster `cluster')
			}
			else if "`pscore'" != "" & `lin_cntrl' == 1 {
				reg `y' `endogbal' `other_controls' `fvars' `pscore', robust
				
			}
			else if "`pscore'" != ""  {
				cap qui: reg `y' `endogbal' `other_controls' `fvars', absorb(`pscore') robust
				
			}
			else{
				qui:	reg `y' `endogbal' `other_controls' `fvars' , robust
				*di in red e(cmdline)
			}
			cap qui: test `instruments'
			
			cap local f = r(F)
			cap local p = r(p)
			cap qui: estadd scalar f = `f'
			cap qui: estadd scalar pvalue = `p'


			cap qui: unique `idvar' if e(sample)
			cap qui: estadd scalar obs=`r(sum)'


			if `mean_offer' == 1 {
				qui: su `endogbal'
				qui: estadd scalar mean_offer=r(mean)

			}
			

			if `non_offered_mean' == 1 {

				local mean_vars

				foreach var of varlist  `endogbal' {
					qui: su `y' if  `var' == 0
					qui: estadd scalar non_`var' = r(mean)
					qui: estadd scalar nonsd_`var'   = r(sd)
					local mean_vars `mean_vars' non_`var' nonsd_`var'
				}

			}

			qui: estimates store e`j'
			
			
			local estimatesstring " `estimatesstring' e`j' "

			local j = `j' + 1
		}


		if "`pointsofsupport'" != "" {

			local points

			foreach var of varlist `pointsofsupport' {
				unique `var'
				estadd scalar  pt_`var' = `r(sum)': e1
				local points `points' pt_`var'

			}
		}

		if `joint' == 1{
			mvreg  baseline_black baseline_hispanic baseline_female baseline_ss_ela baseline_ss_math = `endogbal' `other_controls'
			di in red e(command)
			test `endogbal'
			local Fgeneral = r(F)
			local pgeneral = r(p)

			estadd scalar fgen = `Fgeneral': e1
			estadd scalar pvaluegen = `pgeneral': e1

			print_results  `endogbal', model_names( `estimatesstring' ) hidevars_print( `hidevars' ) ///
			outfile("`outfile'") stat_names( "obs f pvalue fgen pvaluegen `points' mean_offer `mean_vars'") stat_labs(`"obs" "F-test" "p-value" "F-joint" "p-joint" "R2"')
		}
		else{
			print_results  `endogbal', model_names( `estimatesstring' ) hidevars_print( `hidevars' ) ///
			outfile("`outfile'_bal.csv") stat_names( "obs f pvalue `points' mean_offer `mean_vars'") stat_labs(`"obs" "F-test" "p-value" "R2"')
		}
		drop `endogbal'

	}

	if `fs' == 1 {
		//First stage

		* just for esttab
		local endogmod = ""
		foreach var of local instruments{

			gen `var'_fs = `var'
			local endogmod " `endogmod' `var'_fs "
		}
		
		local j = 1

		local estimatesstring ""
		foreach outcome of local outcomes {

		foreach y of local endogvar{

		
			gen `outcome'_`y' = `y' if `outcome' != .
			if "`weight'" != ""{
				qui: reg `outcome'_`y'  `endogmod' `other_controls' `regression_controls' [aw=`weight'], robust

			}
			else if "`pscore'" != "" & `lin_cntrl' == 1 {
				cap qui: reg `outcome'_`y'  `endogmod' `other_controls' `regression_controls' `pscore', robust
				
			}
			else if "`pscore'" != "" {
				qui: reg `outcome'_`y'  `endogmod' `other_controls' `regression_controls', absorb(`pscore') robust
			}
			else if "`cluster'" != "" {
				qui: reg `outcome'_`y'  `endogmod' `other_controls' `regression_controls', vce(cluster `cluster')
			}

			else{
				qui: reg `outcome'_`y'  `endogmod' `other_controls' `regression_controls', robust
			}
			local rsqrd = `e(r2)'
			qui: test  `endogmod'
			local f = r(F)
			local p = r(p)
			qui: estadd scalar f = `f'
			qui: estadd scalar pvalue = `p'

			qui: unique `idvar' if e(sample)
			qui: estadd scalar obs=`r(sum)'

			qui: estimates store e`j'

			local estimatesstring " `estimatesstring' e`j' "

			local j = `j' + 1

		}
		}
		print_results  `endogmod' , model_names( `estimatesstring' )  hidevars_print( `hidevars' ) outfile("`outfile'_fs.csv")	///
									stat_names( " obs f pvalue r2") stat_labs(`"obs" "F test" "p-value" "R2"')

		drop `endogmod'
	}
	 // Reduced form
	if `rf' == 1 {

		* just for esttab
		local endogmod = ""
		foreach var of local instruments{

			gen `var'_rf = `var'
			local endogmod " `endogmod' `var'_rf "
		}
		local estimatesstring ""

		local j = 1
		foreach y of local outcomes{

			if "`weight'" != ""{
				qui: reg `y'  `endogmod' `other_controls' `regression_controls' [aw=`weight'], robust
			}
			else if "`pscore'" != "" & `lin_cntrl' == 1 {
				qui: reg `y'  `endogmod' `other_controls' `regression_controls' `pscore', robust
				
			}
			else if "`pscore'" != "" {
				qui: reg `y'  `endogmod' `other_controls' `regression_controls', absorb(`pscore') robust
			}
			else if "`cluster'" != "" {
				qui: reg `y'  `endogmod' `other_controls' `regression_controls', vce(cluster `cluster')
			}

			else{
				qui: reg `y'  `endogmod' `other_controls' `regression_controls', robust
			}

			qui: test  `endogmod'
			local f = r(F)
			local p = r(p)
			qui: estadd scalar f = `f'
			qui: estadd scalar pvalue = `p'

			qui: unique `idvar' if e(sample)
			qui: estadd scalar obs=`r(sum)'

			qui: estimates store e`j'

			local estimatesstring " `estimatesstring' e`j' "

			if `rf_non_offered_mean' == 1 {
				qui: cap unique `idvar' if e(sample) & `endogmod' == 0
				qui: cap estadd scalar n_non_off = `r(sum)'
			
				
				qui: cap su `y' if `endogmod' == 0
				qui: cap estadd scalar non_off = r(mean)
			}

			di in red "`y'"
			local j = `j' + 1
		}


		print_results  `endogmod' , model_names( `estimatesstring' )  hidevars_print( `hidevars' ) outfile("`outfile'_rf.csv")	///
						stat_names( "obs non_off n_non_off `mean_vars'") stat_labs( `"observations"' `"non-offered mean"' `"non-offered num"' `nonoffered_means' )

		drop `endogmod'

	}

	if `2sls' == 1 {

		//IV estimates
		* just for esttab
		local endogmod = ""
		foreach var of local instruments{

			gen `var'_2s = `var'
			local endogmod " `endogmod' `var'_2s "
		}

		local estimatesstring ""

		local j = 1
		foreach y of local outcomes {
			

			if "`weight'" != ""{
				qui: ivreg2 `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls' [aw=`weight'], robust partial(`other_controls' `regression_controls')
			}
			else if "`cluster'" != "" & "`pscore'" != "" {
				qui: ivreghdfe `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls' , robust absorb(`pscore') cluster(`cluster')
			}
			else if "`pscore'" != ""{
				qui: cap ivreghdfe `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls', robust absorb(`pscore')
			}
			else if "`cluster'" != "" {
				qui: ivreg2 `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls' ,  cluster(`cluster') partial(`other_controls' `regression_controls')
			}
			else if `2sls_residuals' == 1 {
				qui:  ivreg2 `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls', robust  partial(`other_controls' `regression_controls')
				predict res_`y', residuals
			}
			else if "`pscore'" != "" & `lin_cntrl' == 1 {
				qui: cap ivreghdfe `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls' `pscore', robust
			}
			else{
				qui: cap ivreg2 `y' (`endogvar' = `endogmod'  ) `other_controls' `regression_controls', robust partial(`other_controls' `regression_controls')
				
			}
			qui: cap estadd scalar sargan = e(j)
			qui: cap estadd scalar sarganp = e(jp)

			qui: cap unique `idvar' if e(sample)
			qui: cap estadd scalar obs = `r(sum)'
			
			if `non_offered_mean' == 1 {
			
				qui: cap unique `idvar' if e(sample) & `endogmod' == 0
				qui: cap estadd scalar n_non_off = `r(sum)'
			
				
				qui: cap su `y' if `endogmod' == 0
				qui: cap estadd scalar non_off = r(mean)
				

			}

			qui: estimates store e`j'
			
			local estimatesstring " `estimatesstring' e`j' "
	
			qui: if "`model_file'" != "" estimates save "`model_file'_2sls`y'", replace

			local j = `j' + 1

			if `2sls_residuals' == 1 {

				preserve
					keep student_id res_*
					local outfile_resids : subinstr local outfile ".csv" "_2slsresids.dta", all
					sort student_id
					sa "`outfile_resids'", replace
				restore

			}

		}
		print_results  `endogvar' , model_names( `estimatesstring' ) hidevars_print( 0 ) outfile("`outfile'_2sls.csv") ///
		stat_names( "obs non_off n_non_off `mean_vars'") stat_labs( `"observations"' `"non-offered mean"' `"non-offered num"' `nonoffered_means' )
		

		drop `endogmod'

	}


end
